#!/usr/bin/python
import sys
import os
from numpy import *

# Args.
fileins = []
if os.path.isdir(sys.argv[1]):
	for f in os.listdir(sys.argv[1]):
		fileins.append("%s/%s" % (sys.argv[1], f) )
else:
	fileins.append(sys.argv[1])


def N50(numlist):
	# Sort array.
	numlist.sort()
		
	# Compute total sum.
	ts = 0
	for c in numlist:
		ts = ts + c
	target = float(ts) * 0.5
	
	print "Total sequences: %i" % len(numlist)
	print "Cumulitive sum: %i" % ts
	print "Target: %f" % target
	
	# Find first contig for which ts > .5
	cnt = 0.0
	for c in numlist:
		cnt += c
		if cnt >= target:
			return c
	return -1


#assert N50([2, 2, 2, 3, 3, 4, 8, 8]) == 6

# Read in fasta.
lengths = []
for filein in fileins:
	seq = "-1"
	fin = open(filein, "r")
	for line in fin:
		# Skip blanks and comments.
		if len(line) == 0: continue
		if line[0] == "#": continue
		
		# base case.
		if line[0] == ">" and seq == "-1":
			head = line.strip().replace(">","")
			seq = ""
			continue
		
		# Check for ids.
		if line[0] ==  ">": 
			# Save size.
			lengths.append(len(seq))
			
			# reset.
			head = line.strip().replace(">","")
			seq = ""
			continue

		# Filter chars.
		seq += line.strip()
	# Count last sequence.
	lengths.append(len(seq))

# Close file.
fin.close()

print "N50: " + str( N50(lengths) )
